import tensorflow as tf
from tensorflow.keras import datasets, models
import matplotlib.pyplot as plt

# ① 加载MNIST数据集
(x_train, y_train), (x_test, y_test) = datasets.mnist.load_data()

# ② 数据预处理：归一化
x_train = x_train.astype('float32') / 255.0
x_test = x_test.astype('float32') / 255.0

# ③ 调整形状为 [样本数, 28, 28, 1]
x_train = x_train.reshape(-1, 28, 28, 1)
x_test = x_test.reshape(-1, 28, 28, 1)

# ④ 标签one-hot编码
y_train = tf.keras.utils.to_categorical(y_train, 10)
y_test = tf.keras.utils.to_categorical(y_test, 10)

# ⑤ 构建模型
model = models.Sequential([
    tf.keras.layers.Conv2D(64, (3, 3), strides=(1, 1), activation='relu',
                           padding='same', input_shape=(28, 28, 1)),
    tf.keras.layers.MaxPooling2D((2, 2)),
    tf.keras.layers.Conv2D(32, (3, 3), activation='relu', padding='same'),
    tf.keras.layers.Conv2D(64, (3, 3), activation='relu', padding='same'),
    tf.keras.layers.MaxPooling2D((2, 2)),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(128, activation='relu'),
    tf.keras.layers.Dense(10, activation='softmax')
])

# ⑥ 打印模型摘要
model.summary()

# ⑦ 编译模型
model.compile(optimizer='adam',
              loss='categorical_crossentropy',
              metrics=['accuracy'])

# ⑧ 训练模型（序号⑧的代码不完整，补充batch_size和epochs）
history = model.fit(x_train, y_train,
                    batch_size=64,
                    epochs=5,
                    validation_data=(x_test, y_test))

# ⑨ 评估模型
test_loss, test_acc = model.evaluate(x_test, y_test, verbose=2)
print('\nTest accuracy:', test_acc)

# ⑩ 可视化训练过程（可选）
plt.plot(history.history['accuracy'], label='accuracy')
plt.plot(history.history['val_accuracy'], label='val_accuracy')
plt.xlabel('Epoch')
plt.ylabel('Accuracy')
plt.ylim([0, 1])
plt.legend(loc='lower right')
plt.show()

# ⑪ 构建新模型（对比实验）
model2 = models.Sequential([
    tf.keras.layers.Conv2D(32, (3, 3), activation='relu',
                           input_shape=(28, 28, 1)),
    tf.keras.layers.MaxPooling2D((2, 2)),
    tf.keras.layers.Conv2D(64, (3, 3), activation='relu'),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(10, activation='softmax')
])

# ⑫ 编译新模型
model2.compile(optimizer='adam',
               loss='sparse_categorical_crossentropy',  # 注意这里使用sparse版本
               metrics=['accuracy'])

# ⑬ 训练新模型（使用原始标签而非one-hot）
history2 = model2.fit(x_train.reshape(-1, 28, 28, 1), y_train.argmax(axis=1),
                      batch_size=64,
                      epochs=3,
                      validation_data=(x_test.reshape(-1, 28, 28, 1), y_test.argmax(axis=1)))

# ⑭ 评估新模型
test_loss2, test_acc2 = model2.evaluate(x_test.reshape(-1, 28, 28, 1), y_test.argmax(axis=1),
                                        verbose=2)
print('\nTest accuracy (model2):', test_acc2)

# ⑮ 可视化对比（可选）
plt.plot(history.history['val_accuracy'], label='model1')
plt.plot(history2.history['val_accuracy'], label='model2')
plt.xlabel('Epoch')
plt.ylabel('Validation Accuracy')
plt.legend()
plt.show()